import numpy as np
import pandas as pd
from lazypredict.Supervised import LazyClassifier
from sklearn.model_selection import train_test_split
from util import createXY
import faiss
import logging
import joblib
import gradio as gr
import cv2
import argparse
# 配置logging, 确保能够打印正在运行的函数名
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
#禁用joblib的日志输出
logging.getLogger('joblib').setLevel(logging.WARNING)

def get_args():
    parser = argparse.ArgumentParser(description='使用CPU或GPU训练模型。') # 创建命令行参数解析器对象
    parser.add_argument('-m', '--mode', type=str, required=True, choices=['cpu', 'gpu'], help='选择训练模式：CPU或GPU。')
    parser.add_argument('-f', '--feature', type=str, required=True, choices=['flat', 'vgg'], help='选择特征提取方法：flat或vgg。')
    parser.add_argument('-l', '--library', type=str, required=True, choices=['sklearn', 'faiss'], help='选择使用的库：sklearn或faiss。')
    args = parser.parse_args()
    return args

args = get_args()

# 根据mode初始化FAISS所需的资源
res = faiss.StandardGpuResources() if args.mode == 'gpu' else None


# 加载数据
X, y = createXY(train_folder="data/train", dest_folder=".",method=args.feature)
X = np.array(X).astype('float32')
faiss.normalize_L2(X)  # 对数据进行L2归一化
y = np.array(y)
logging.info("数据加载和预处理完成。")

# 数据集分割为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
logging.info("数据集划分为训练集和测试集。")
print(len(X_train), len(y_train)) 

# 定义模型评估的数据
clf = LazyClassifier()
result, _ = clf.fit(X_train, X_test, y_train, y_test)
print(result)
# 保存最好的模型
best_model_name = result['F1 Score'].idxmax()
print("\nF1 Score best model:", best_model_name)
best_model = clf.models[best_model_name]
result = best_model.predict(X_test)
print(f"use {best_model_name} predict X_test:\n{result}")
joblib.dump(best_model, open('best_lazypredict_model.pkl', 'wb'))

# 加载模型
model = joblib.load('best_lazypredict_model.pkl')

# 定义分类函数
def predict_image(img):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 将图像从BGR格式转换为RGB格式
    img = cv2.resize(img, (32, 32))  # 调整图像大小到32x32
    img = img.reshape(1,-1)  # 将图像展平
    prediction = model.predict(img)
    return "Dog" if prediction[0]==1 else "Cat"

# 创建 Gradio 界面
iface = gr.Interface(
    inputs="image",
    outputs="label",
    fn=predict_image
)

# 启动 Gradio 界面
iface.launch()  # 启动界面